import ast
from _ast import Return, Break, FunctionDef, ClassDef, Name
from ast import walk, NodeVisitor, NodeTransformer

import typing


class NameReplacer(NodeTransformer):
    def __init__(self, rules: typing.Dict[str, str]):
        self.rules = rules

    def visit_Name(self, node: Name) -> Name:
        if node.id in self.rules:
            return Name(id=self.rules[node.id], ctx=node.ctx)
        else:
            return node


# def query_first_node_by_type(tree: typing.Union[ast.AST, typing.List[ast.AST]],
#                              node_type: typing.Type[ast.AST]):
#     assert issubclass(node_type, ast.AST)
#     if isinstance(tree, ast.AST):
#         tree = [tree]
#     for expr in tree:
#         for item in ast.walk(expr):
#             if isinstance(item, node_type):
#                 # all_nodes.append(item)
#                 return item
#     return None


# def query_first_node_by_type(tree: typing.Union[ast.AST, typing.List[ast.AST]],
#                              node_type: typing.Type[ast.AST]):
#     assert issubclass(node_type, ast.AST)
#     if isinstance(tree, ast.AST):
#         tree = [tree]
#     for expr in tree:
#         for item in ast.walk(expr):
#             if isinstance(item, node_type):
#                 # all_nodes.append(item)
#                 return item
#     return None


# def query_all_by_type(tree: ast.AST, node_type: typing.Type[ast.AST]):
#     assert issubclass(node_type, ast.AST)
#     all_nodes: typing.List[ast.AST] = []
#     for item in ast.walk(tree):
#         if isinstance(item, node_type):
#             all_nodes.append(item)
#     return all_nodes


# def has_return(tree: ast.AST):
#     return query_first_node_by_type(tree, Return) is not None


# def has_jump_out(tree: ast.AST):
#     has_return_node = query_first_node_by_type(tree, Return) is not None
#     has_break_node = query_first_node_by_type(tree, Break) is not None
#     return not (has_return_node or has_break_node)


def replace_all_names(tree: ast.AST, rules):
    NameReplacer(rules).visit(tree)

# def has_generic_annotation(tree: typing.Union[ast.FunctionDef, ast.ClassDef], generic_names: typing.List[str]):
#
